﻿// Copyright © .NET Foundation and Contributors. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.

namespace PInvoke;

using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Security.Cryptography;
using System.Text;
using Microsoft.CodeAnalysis;
using Microsoft.CodeAnalysis.CSharp;
using Microsoft.CodeAnalysis.CSharp.Syntax;
using static Microsoft.CodeAnalysis.CSharp.SyntaxFactory;

[Generator]
public class SourceGenerator : ISourceGenerator
{
    /// <summary>
    /// A "generated by tool" comment string with environment/os-normalized newlines.
    /// </summary>
    public static readonly SyntaxTriviaList GeneratedByAToolPreamble = ParseLeadingTrivia(@"// ------------------------------------------------------------------------------
// <auto-generated>
//     This code was generated by a tool.
//
//     Changes to this file may cause incorrect behavior and will be lost if
//     the code is regenerated.
// </auto-generated>
// ------------------------------------------------------------------------------
".Replace("\r\n", "\n").Replace("\n", Environment.NewLine)); // normalize regardless of git checkout policy

    public void Execute(GeneratorExecutionContext context)
    {
        try
        {
            using HashAlgorithm hash = HashAlgorithm.Create("SHA1");

            foreach (SyntaxTree syntaxTree in context.Compilation.SyntaxTrees)
            {
                SemanticModel semanticModel = context.Compilation.GetSemanticModel(syntaxTree);
                CompilationUnitSyntax? root = syntaxTree.GetRoot(context.CancellationToken) as CompilationUnitSyntax;
                if (root is null)
                {
                    continue;
                }

                SyntaxList<MemberDeclarationSyntax> additionalMembers = default;
                foreach (TypeDeclarationSyntax typeDeclaration in root.DescendantNodes().OfType<TypeDeclarationSyntax>())
                {
                    INamedTypeSymbol? applyTo = semanticModel.GetDeclaredSymbol(typeDeclaration, context.CancellationToken);
                    if (applyTo is null)
                    {
                        continue;
                    }

                    foreach (AttributeData att in applyTo.GetAttributes())
                    {
                        IGenerator? generator = att.AttributeClass?.Name switch
                        {
                            nameof(OfferFriendlyOverloadsAttribute) => new OfferFriendlyOverloadsGenerator(),
                            nameof(OfferIntPtrPropertyAccessorsAttribute) => new OfferIntPtrPropertyAccessorsGenerator(),
                            _ => null,
                        };

                        if (generator is object)
                        {
                            TransformationContext transformationContext = new TransformationContext
                            {
                                Compilation = context.Compilation,
                                ProcessingNode = typeDeclaration,
                            };
                            SyntaxList<MemberDeclarationSyntax> generatedMembers = generator.Generate(transformationContext, context.CancellationToken);
                            if (generatedMembers.Any(m => m.DescendantNodesAndSelf().OfType<TypeDeclarationSyntax>().Any(t => t.Members.Any())))
                            {
                                SyntaxList<MemberDeclarationSyntax> wrappedMembers = typeDeclaration.Ancestors().Aggregate(generatedMembers, WrapInAncestor);
                                additionalMembers = additionalMembers.AddRange(wrappedMembers);
                            }
                        }
                    }
                }

                if (additionalMembers.Count > 0)
                {
                    var emittedExterns = root.Externs.Select(x => x.WithoutTrivia()).ToImmutableArray();
                    var emittedUsings = root.Usings.Select(x => x.WithoutTrivia()).ToImmutableArray();

                    CompilationUnitSyntax generatedFile = CompilationUnit()
                        .WithExterns(List(emittedExterns))
                        .WithUsings(List(emittedUsings))
                        .WithMembers(additionalMembers)
                        .WithLeadingTrivia(GeneratedByAToolPreamble)
                        .NormalizeWhitespace();

                    string sourceHash = Convert.ToBase64String(hash.ComputeHash(Encoding.UTF8.GetBytes(syntaxTree.FilePath)), 0, 6);
                    string hintName = Path.GetFileNameWithoutExtension(syntaxTree.FilePath) + $".{sourceHash}.friendly.cs";
                    context.AddSource(hintName.Replace('/', '-').Replace('+', '_'), generatedFile.ToFullString());
                }
            }
        }
        catch (Exception ex) when (LaunchDebugger(ex))
        {
            throw;
        }

        static bool LaunchDebugger(Exception ex)
        {
            if (ex is OperationCanceledException)
            {
                return false;
            }

#if DEBUG
            Debugger.Launch();
#endif
            return false;
        }
    }

    public void Initialize(GeneratorInitializationContext context)
    {
    }

    private static SyntaxList<MemberDeclarationSyntax> WrapInAncestor(SyntaxList<MemberDeclarationSyntax> generatedMembers, SyntaxNode ancestor)
    {
        switch (ancestor)
        {
            case NamespaceDeclarationSyntax ancestorNamespace:
                generatedMembers = SyntaxFactory.SingletonList<MemberDeclarationSyntax>(
                    CopyAsAncestor(ancestorNamespace)
                    .WithMembers(generatedMembers));
                break;
            case ClassDeclarationSyntax nestingClass:
                generatedMembers = SyntaxFactory.SingletonList<MemberDeclarationSyntax>(
                    CopyAsAncestor(nestingClass)
                    .WithMembers(generatedMembers));
                break;
            case StructDeclarationSyntax nestingStruct:
                generatedMembers = SyntaxFactory.SingletonList<MemberDeclarationSyntax>(
                    CopyAsAncestor(nestingStruct)
                    .WithMembers(generatedMembers));
                break;
        }

        return generatedMembers;
    }

    private static NamespaceDeclarationSyntax CopyAsAncestor(NamespaceDeclarationSyntax syntax)
    {
        return SyntaxFactory.NamespaceDeclaration(syntax.Name.WithoutTrivia())
            .WithExterns(SyntaxFactory.List(syntax.Externs.Select(x => x.WithoutTrivia())))
            .WithUsings(SyntaxFactory.List(syntax.Usings.Select(x => x.WithoutTrivia())));
    }

    private static ClassDeclarationSyntax CopyAsAncestor(ClassDeclarationSyntax syntax)
    {
        return SyntaxFactory.ClassDeclaration(syntax.Identifier.WithoutTrivia())
            .WithModifiers(SyntaxFactory.TokenList(syntax.Modifiers.Select(x => x.WithoutTrivia())))
            .WithTypeParameterList(syntax.TypeParameterList);
    }

    private static StructDeclarationSyntax CopyAsAncestor(StructDeclarationSyntax syntax)
    {
        return SyntaxFactory.StructDeclaration(syntax.Identifier.WithoutTrivia())
            .WithModifiers(SyntaxFactory.TokenList(syntax.Modifiers.Select(x => x.WithoutTrivia())))
            .WithTypeParameterList(syntax.TypeParameterList);
    }
}
